
from typing import List
import os
import datetime
from PIL import Image
import numpy as np
# import folder_paths
import tensorflow as tf
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.outputs import OutputKeys
from alioffline_utils_img import *
from alioffline_utils_mask import png_to_mask

# comfy_path = os.path.dirname(folder_paths.__file__)
current_directory = os.path.dirname(__file__)
comfy_path =  os.path.abspath(os.path.join(current_directory, '..', '..'))
models_path = os.path.join(comfy_path, "models","modelscope")
# model_path = os.path.join(models_path, "demo","cv_ddcolor_image-colorization")
  
class AliOffline_Ddcolor:
   
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
              {   
                "image": ('IMAGE', {}),
              }
            }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    OUTPUT_NODE = True
    FUNCTION = "sample"
    CATEGORY = "CXH"

    def sample(self,image):
        os.environ['MODELSCOPE_CACHE']=models_path
        #res_images = []
        tf.compat.v1.disable_eager_execution()
        sess = tf.compat.v1.Session()
        model = pipeline(
            Tasks.image_colorization, model="damo/cv_ddcolor_image-colorization"
        )
      #for item in image:
        result = model(tensor2pil(item))
        ndarray = cv2.cvtColor(result[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGBA)
        out_image = array2image(ndarray)
        out_tensor = pil2tensor(out_image)
        #  res_images.extend(out_tensor)
        sess.close()
        return (out_tensor,)